
from typing import Dict, List
import numpy as np
from pylab import mpl, plt
from matplotlib import rcParams

from obspy.geodetics.base import gps2dist_azimuth

deg2rad = 1/180*np.pi
rad2deg = 1/np.pi*180

rcParams.update({
    'font.family': ['sans-serif'],
    'font.sans-serif': [
        'Noto Sans CJK JP',  # 中文字体（思源黑体）
        'Nimbus Sans',       # 英文字体
        'DejaVu Sans',       # 备用英文字体
        'Arial',             # 备用
        'sans-serif'         # 最终备用
    ],
    'font.size': 10,
    'axes.titlesize': 12,
    'axes.labelsize': 11,
    'xtick.labelsize': 9,
    'ytick.labelsize': 9,
    'legend.fontsize': 9,
    
    # 数学字体设置（重要！）
    'mathtext.fontset': 'stix',      # 使用STIX数学字体
    'mathtext.default': 'regular',   # 常规样式
})


def load_loc(filename='./data/XA.Loc.201804.txt', name_key='name', lat_key = 'lat_obs', lon_key = 'lon_obs', sn_key=None) -> Dict:
    stations = {}
    if '.csv' in filename:
        import pandas as pd
        loc = pd.DataFrame(pd.read_csv(filename))
        lats = list(loc[lat_key])
        lons = list(loc[lon_key])
        names = list(loc[name_key])
        if sn_key is not None:
            sn = list(loc[sn_key])
            for i in range(len(names)):
                lat = float(lats[i])
                lon = float(lons[i])
                stations[names[i]] = [lat, lon, sn[i]]
        else:
            for i in range(len(names)):
                lat = float(lats[i])
                lon = float(lons[i])
                stations[names[i]] = [lat, lon]
                
        return stations
    
    with open(filename,'r') as file:
        inf = file.readlines()
    for stat_inf in inf:
        if '2019' in filename:
            name = stat_inf[0:5]
            lat = float(stat_inf[10:21])
            lon = float(stat_inf[25:37])
        elif '2018' in filename:
            name = stat_inf.split('\t')[0]
            lon = float(stat_inf.split('\t')[1])
            lat = float(stat_inf.split('\t')[2])
        if '201805' in filename:
            name = stat_inf[0:5]
        
        stations[name] = [lat, lon]

    return stations

def plot_stations_loc(stations : Dict,names2show=[], points2show=[], 
                      KM=True, special_names=[], c='',
                      map_style = 2,bounds=None,zoom=14,ax=None, dx=0.02,dy=0.01,
                      filename=None,
                      ) -> None:
    '''
    绘制地震台站位置
    stations : 台站位置信息
    names2show : 显示的台站名称
    points2show : 显示的点名称
    KM : 是否以KM为单位显示
    special_names : 特殊台站名称
    map_style : 地图样式
    bounds : 地图边界
    zoom : 地图缩放比例
    filename : 保存图片名称
    map_style : 地图样式
    ax : 绘制的轴
    dx : 经度间隔
    dy : 纬度间隔
    filename : 保存图片名称
    '''
    import transbigdata as tbd

    rcParams['font.family'] = 'Arial'
    rcParams['font.size'] = '10'

    if ax is None:
        fig = plt.figure(figsize=[5,4], dpi=600)
        ax = plt.subplot(111)

    if bounds is None:
        # bounds = [115.16,38.57,115.36,38.68]
        bounds = [115.22,38.60,115.26,38.63]
        # bounds = [115.66,39.06,115.92,39.91]
        # bounds = [115.77,39.00,115.92,39.14]
    #创建图框
    plt.sca(ax)
    #添加地图底图
    tbd.plot_map(plt,bounds,zoom = zoom,style = map_style)
    # tbd.plotscale(ax,bounds = bounds,textsize = 5,compasssize = 1,accuracy = 300,unit='m',rect = [0.06,0.9],zorder = 15)
    color='k' if map_style==2 else 'red'
    lats, lons = [],[]
    for stat_name in stations.keys():
        lat, lon= stations[stat_name][:2]
        
        if stat_name in names2show:
            plt.text(lon, lat, stat_name, color='tab:red', fontdict={'fontsize':10}, rotation=0)
            plt.scatter(lon, lat,marker='^', s=30, color='tab:red', edgecolor='none')
        elif stat_name in points2show:
            plt.scatter(lon, lat,marker='^', s=25, color='blue', edgecolor='none')
        if stat_name in names2show or stat_name in points2show:
            lats.append(lat)
            lons.append(lon)
    
    # xs,xe = int(min(lons)/dx)*dx,int(max(lons)/dx)*dx
    # ys,ye = int(min(lats)/dy)*dy,int(max(lats)/dy+1)*dy
    ys,ye = bounds[1], bounds[3]
    xs,xe = bounds[0], bounds[2]
    ax.set_xlim([xs,xe])
    ax.set_ylim([ys,ye])
    xticks = np.arange(xs,xe,dx)
    xtickslabel = ['{:3.2f}'.format(i) for i in xticks]
    ax.set_xticks(xticks)
    ax.set_xticklabels(xtickslabel)
    yticks = np.arange(ys,ye+dy,dy)
    ytickslabel = ['{:2.2f}'.format(i) for i in yticks]
    ax.set_yticks(yticks)
    ax.set_yticklabels(ytickslabel)
    
    
    if filename is not None:
        plt.xlabel('longitude')
        plt.ylabel('latitude')
        plt.tight_layout()
        plt.savefig(filename)
    return ax


def plot_stations_loc_H1(stations : Dict,names2show=[], KM=True, filename=None, map_style = 2,ax=None, dx=0.02,dy=0.01, special_names=[], c='') -> None:
    rcParams['font.family'] = 'Times New Roman'
    rcParams['font.size'] = '8'
    import transbigdata as tbd

    if ax is None:
        fig = plt.figure(figsize=[4,3], dpi=600)
        ax = plt.subplot(111)

    # bounds = [115.16,38.57,115.36,38.68]
    bounds = [115.22,38.60,115.26,38.63]
    # bounds = [115.66,39.06,115.92,39.91]
    # bounds = [115.77,39.00,115.92,39.14]
    #创建图框
    plt.sca(ax)
    #添加地图底图
    tbd.plot_map(plt,bounds,zoom = 14,style = map_style)
    # tbd.plotscale(ax,bounds = bounds,textsize = 5,compasssize = 1,accuracy = 300,unit='m',rect = [0.06,0.9],zorder = 15)
    color='k' if map_style==2 else 'red'
    lats, lons = [],[]
    for stat_name in stations.keys():
        lat, lon= stations[stat_name]
        
        if stat_name in names2show:
            if 'DZ06' in stat_name:
                plt.text(lon, lat-0.004, stat_name, color='red', fontdict={'fontsize':7}, rotation=-45)
                plt.scatter(lon, lat,marker='^', s=25, color='red', edgecolor='none')
            elif 'CF' in stat_name:
                plt.text(lon, lat-0.003, stat_name, color='orange', fontdict={'fontsize':7}, rotation=-45)
                plt.scatter(lon, lat,marker='^', s=20, color='orange', edgecolor='none')
            else:
                plt.text(lon, lat, stat_name, color='orange', fontdict={'fontsize':7}, rotation=45)
                plt.scatter(lon, lat,marker='^', s=30, color='orange', edgecolor='none')
            lats.append(lat)
            lons.append(lon)
    
    xs,xe = int(min(lons)/dx)*dx,int(max(lons)/dx+1)*dx
    ys,ye = int(min(lats)/dy)*dy,int(max(lats)/dy+1)*dy
    ys,ye = bounds[1], bounds[3]
    xs,xe = bounds[0], bounds[2]

    xticks = np.arange(xs,xe+dx,dx)
    xtickslabel = ['{:3.2f}'.format(i) for i in xticks]
    ax.set_xticks(xticks)
    ax.set_xticklabels(xtickslabel)
    yticks = np.arange(ys,ye+dy,dy)
    ytickslabel = ['{:2.2f}'.format(i) for i in yticks]
    ax.set_yticks(yticks)
    ax.set_yticklabels(ytickslabel)

    ax.set_xlim([xs,xe])
    ax.set_ylim([ys,ye])
    
    
    if filename is not None:
        plt.xlabel('longitude')
        plt.ylabel('latitude')
        plt.tight_layout()
        plt.savefig(filename)
        print(filename)
    return ax


def get_distance(name1=None, name2=None, loc1=None, loc2=None, S2N=True,s_info=None,EPISON=5, AZ_HSR=55):
    
    if (name1 is None and loc1 is None) or (name2 is None and loc2 is None):
        return None
    if (name1 is not None or name2 is not None) and s_info is None:
        return None
    lat1, lon1 = loc1 if name1 is None else s_info[name1][:2]
    lat2, lon2 = loc2 if name2 is None else s_info[name2][:2]
    distance,AZ,BAZ = gps2dist_azimuth(lat1, lon1, lat2, lon2)
    if S2N and AZ>AZ_HSR+90+EPISON and AZ<AZ_HSR+270-EPISON:
        distance = - distance
    
    return distance,AZ,BAZ

def AZ2theta(AZ,AZ_HSR=55):
    a = np.cos(deg2rad*(AZ-AZ_HSR))

    theta = np.arccos(a)*rad2deg
    
    return theta


def sort_data_by_distance(distances, DATA=None, NAMES=None):
    '''
    根据距离对数据、名称等进行排序
    distances: 距离列表
    DATA: 数据列表
    NAMES: 名称列表

    return distances_sorted,data_sorted,names_sorted
    '''
    ns = len(distances)
    if NAMES is not None:
        assert ns==len(NAMES)
        names=NAMES
    else:
        names=['0']*ns
    if DATA is not None:
        assert ns==DATA.shape[0]
        data=DATA
    else:
        data=np.zeros([ns,1])
    if NAMES is None and DATA is None:
        raise ValueError('At least one para in DATA and NAMES is not None')

    list_data = []
    for i in range(ns):
        trace = data[i,:]
        dist_i = distances[i]
        name_i = names[i]
        list_data.append([dist_i,trace,name_i])
    list_data.sort(key=lambda x:x[0])

    distances_sorted=[]
    data_sorted=np.zeros_like(data)
    names_sorted=[]
    for i in range(ns):
        distances_sorted.append(list_data[i][0])
        data_sorted[i,:]=list_data[i][1]
        names_sorted.append(list_data[i][2])
    if NAMES is None:
        return distances_sorted,data_sorted
    elif DATA is None:
        return distances_sorted,names_sorted
    else:
        return distances_sorted,data_sorted,names_sorted

def get_stats_within_distance(s_info, base, r0=-1000, r1=1000):
    '''
    获取名称name的站内距离小于r0的所有台站
    s_info: 站信息
    base: 基站名称
    r0: 最小距离
    r1: 最大距离
    return: 站名称列表
    '''
    stats_in_range = []
    
    # 获取基站的位置信息
    if base not in s_info:
        raise ValueError(f"基站 {base} 不在台站信息中")
    
    # 遍历所有台站
    for stat_name in s_info.keys():
        # 计算与其他台站的距离
        distance, _, _ = get_distance(name1=base, name2=stat_name, s_info=s_info,S2N=True)

        # 判断是否在指定范围内
        if r0 <= distance <= r1:
            stats_in_range.append(stat_name)
    
    return stats_in_range


def get_stats_within_aCircle(s_info, base, r=1000):
    '''
    获取名称name的站内距离小于r的所有台站
    s_info: 站信息
    base: 基站名称
    r: 距离，单位米
    return: 站名称列表
    '''
    stats_in_range = []
    
    # 获取基站的位置信息
    if base not in s_info:
        raise ValueError(f"基站 {base} 不在台站信息中")
    
    # 遍历所有台站
    for stat_name in s_info.keys():
        # 计算与其他台站的距离
        distance, _, _ = get_distance(name1=base, name2=stat_name, s_info=s_info,S2N=False)

        # 判断是否在指定范围内
        if distance <= r:
            stats_in_range.append(stat_name)
    
    return stats_in_range


def test():
    '''
    绘制一个
    119度48分
    44度16分
    119度54分
    44度20分
    的地图
    '''
    # 创建虚拟台站数据用于测试
    stations = {
        'STA01': [44.2667, 119.8],  # 44度16分 = 44.2667度
        'STA02': [44.3333, 119.9],  # 44度20分 = 44.3333度, 119度54分 = 119.9度
        'STA03': [44.3, 119.85],    # 中间位置
    }
    
    # 定义地图边界 [西, 南, 东, 北]
    # 119度48分 = 119.8度, 119度54分 = 119.9度
    # 44度16分 = 44.2667度, 44度20分 = 44.3333度
    bounds = [119+48/60, 44.23, 119+54/60, 44.33]
    
    # 创建绘图
    fig = plt.figure(figsize=[6, 4], dpi=600)
    ax = plt.subplot(111)
    
    # 绘制台站位置
    plot_stations_loc(
        stations={},
        names2show=[],
        bounds=bounds,
        zoom=14,
        ax=ax,
        map_style=3
    )
    
    # 设置标题和标签
    plt.title('Test Map')
    plt.xlabel('Longitude')
    plt.ylabel('Latitude')
    
    # 显示图形
    plt.tight_layout()
    fig.savefig('loc_test_map.png')
    fig.savefig('loc_test_map.pdf')
    
if __name__ == '__main__':
    test()